Split reverse-mode thunk for the IIP Enzyme rule (alternative to #63)#64
Merged
ChrisRackauckas merged 1 commit intoJul 5, 2026
Conversation
…ciML#63) Reimplements the Const-return (IIP) reverse rule by delegating to Enzyme's own split reverse-mode `autodiff_thunk` on the *unwrapped* function, rather than snapshotting arguments and re-running `autodiff(Reverse, …)` in the reverse pass (the approach in the snapshot PR): - `augmented_primal` builds `autodiff_thunk(ReverseSplitModified(…, overwritten))` and runs its forward pass; the `ModifiedBetween == overwritten` flags make Enzyme's own tape cache any argument the caller mutates before the reverse pass. It stashes `(tape, reverse_thunk)`. - `reverse` runs the stashed reverse thunk; `_revslot` rebuilds its per-argument return exactly typed (so `Active` args are correct, not zeroed / union-typed). This is the more Enzyme-native design and passes every isolated test (mutation, multi-step integrator, mixed Duplicated/Active). It is offered as an alternative to SciML#63. KNOWN LIMITATION: on the full nested `EnsembleProblem` adjoint (SciMLSensitivity #1424) the nested `autodiff_thunk` inside `augmented_primal` re-triggers the Julia-1.11 `_dispatch_ensemble_solve` GC-root-rewrite segfault in Enzyme's `post_optimize!` (EnzymeAD/Enzyme.jl #3175), which the lightweight snapshot rule in SciML#63 avoids. So this should not merge ahead of SciML#63 until #3175 is fixed upstream; it is the target design once it is. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Enzyme-native version of the IIP reverse-rule fix. Instead of snapshotting arguments and re-running
autodiff(Reverse, …)in the reverse pass (#63), this delegates to Enzyme's own split reverse-modeautodiff_thunkon the unwrapped function:augmented_primal→autodiff_thunk(ReverseSplitModified(mode, Val(overwritten(config))), Const{typeof(f)}, Const, argtypes...), runs the forward pass, stashes(tape, reverse_thunk). ForwardingoverwrittenasModifiedBetweenmakes Enzyme's own tape cache the mutated arguments (the ODE-integratoru), so there's no manual snapshot/restore.reverseruns the stashed reverse thunk;_revslotrebuilds its per-arg return exact-typed soActiveargs are correct (not zeroed / union-typed →ReverseRuleReturnError).Passes the whole
test/Enzymesuite, including the mutation, multi-step-integrator, and mixedDuplicated/Activeregression tests.On the full nested
EnsembleProblemadjoint (SciML/SciMLSensitivity.jl#1424), the nestedautodiff_thunkinsideaugmented_primalre-triggers the Julia-1.11_dispatch_ensemble_solveGC-root-rewrite segfault in Enzyme'spost_optimize!— the #3175 crash that the lightweight snapshot rule in #63 avoids. Verified: identical script,AutoSpecializeensemble adjoint →#63= exit 0 (no segfault), this branch = SIGSEGV inoptimize.jl:485.So #63 is the mergeable fix today; this is the target design once EnzymeAD/Enzyme.jl#3175 is resolved (a SciML-free reduction of #3175 wasn't found — a self-contained split-thunk custom rule over a 100-trajectory × 120-array aggregate compiles fine; the crash needs the real
_dispatch_ensemble_solveIR).🤖 Generated with Claude Code